import ast
import operator
from pyomo.environ import *
from functools import partial
from numbers import Number
from string import Template
from pyomo.core.expr.numeric_expr import ExpressionBase
from pyomo.core.base.var import Var,AbstractScalarVar,ScalarVar
from pyomo.core.base.param import Param

def myabs(num):
    try:
        if isinstance(num,(Number,AbstractScalarVar,ExpressionBase)) :
            return abs(num)
        else:
            return list(map(abs,num))
    except Exception as e :
        raise e.with_traceback(e.__traceback__)


def mysum(num):
    try:
        if isinstance(num,(Var,Param)):
            return quicksum(num[i] for i in num)
        if isinstance(num,list):
            result = 0
            for item in num:
                result += item
            return result
        else:
            raise RuntimeError(f" unsupport type of {type(num)=} of summation")
    except Exception as e:
        raise e.with_traceback(e.__traceback__)
    
def mypow(num,exp=2):
    try:
        if isinstance(num,(Number,ScalarVar,AbstractScalarVar)) :
            return pow(num,exp=exp)
        else:
            return list(map(lambda x:pow(x,exp),num))
    except Exception as e:
        raise e.with_traceback(e.__traceback__)


func_map = {
    'sum':mysum,
    'abs':myabs,
    'pow':mypow,
}



def evaluate_expr(expr, model):
    def _eval(node):
        if isinstance(node, ast.BinOp):
            left = _eval(node.left)
            right = _eval(node.right)
            op = _get_operator(node.op)
            result = do_calculation(left, right,op)
            return result
        elif isinstance(node, ast.Num):
            return node.n 
        elif isinstance(node,ast.Attribute):
            return getattr(model, node.attr)
        elif isinstance(node, ast.Call):
            arg_result = [_eval(arg) for arg in node.args]
            return func_map[node.func.id](*arg_result)


    def do_calculation(left, right, op):
        _left = left 
        _right = right 
        l_index_list = []
        r_index_list = []
        if isinstance(left,(Var,Param,list)) and isinstance(right,(Var,Param,list)):
            l_index_list = left.index_set() if not isinstance(left,list) else list(range(len(left)))
            r_index_list = right.index_set() if not isinstance(right,list) else list(range(len(right)))
        elif isinstance(left,(Number,ExpressionBase)) and (isinstance(right,(Var,list)) ):
            _left = [left ] * len(right)
            l_index_list = list(range(0,len(_left)))
            r_index_list = list(range(len(right))) if isinstance(right,list) else right.index_set()
        elif isinstance(right,(Number,ExpressionBase)) and isinstance(left,(Var,list)):
            _right = [right] * len(left)
            r_index_list = list(range(0,len(_right)))
            l_index_list = list(range(len(left))) if isinstance(left,list) else left.index_set()
        elif isinstance(left,list) and isinstance(right,list):
            l_index_list = list(range(0,len(left)))
            r_index_list = list(range(0,len(right)))
        elif isinstance(left,(Number,ExpressionBase)) and isinstance(right,(Number,ExpressionBase)):
            return op(left,right)
        else:
            raise TypeError(f"unsupprt type of {type(left)=} {type(right)=}")

        assert len(l_index_list) == len(r_index_list),"左右操作数，长度不同"

        if op == operator.add:
            return [ _left[l] + _right[r] for l,r in zip(l_index_list,r_index_list)]
        elif op == operator.sub:
            return [ _left[l] - _right[r] for l,r in zip(l_index_list,r_index_list)]
        elif op == operator.mul:
            return [ _left[l] * _right[r] for l,r in zip(l_index_list,r_index_list)]
        elif op == operator.truediv:
            return [ _left[l] / _right[r] for l,r in zip(l_index_list,r_index_list)]
        else:
            raise TypeError(op)

    def _get_operator(op):
        if isinstance(op, ast.Add):
            return operator.add
        elif isinstance(op, ast.Sub):
            return operator.sub
        elif isinstance(op, ast.Mult):
            return operator.mul
        elif isinstance(op, ast.Div):
            return operator.truediv
        else:
            raise TypeError(op)

    tree = ast.parse(expr, mode='eval')
    return _eval(tree.body)

model = ConcreteModel()
model.a = Var([1,2,3], within=NonNegativeReals,bounds=(0, 6),initialize=1)
model.b = Var([1,2,3], within=NonNegativeReals,bounds=(0, 6),initialize=1)
model.c = Var([1,2,3], within=NonNegativeReals,bounds=(0, 6),initialize=1)
model.d = Var([1,2,3], within=NonNegativeReals,bounds=(0, 6),initialize=1)

# expr = Template("""( sum($a) + sum($b) +1 ) / (sum($c) + sum($d) + 2)""")
# obj = """sum(($a + $b) / pow(( abs($c - $d ) +1),2) )"""
obj = """sum($a / sum($a) * abs(($a * (1-$b) - $c * $d)))"""
expr = Template(obj)
variable_mapping = {
    "a":"model.a",
    "b":"model.b",
    "c":"model.c",
    "d":"model.d"
}

#
# pyomo download-extensions
# pyomo build-extensions
expr = expr.substitute(variable_mapping)

fn = partial(evaluate_expr,expr)
model.obj = Objective(rule=fn) #
print(model.obj.expr)
opt = SolverFactory('cyipopt')
opt.solve(model,tee=False,options={'max_iter':500,'print_level':5},halt_on_evaluation_error=False)
model.display()

# result = evaluate_expr(expr,model)
# # print(f"Final result: {result}")
print(expr) # py38 dont have indent 
print(ast.dump(ast.parse(expr,mode='eval')))

# print(model.a[1] + model.a[2])